/*
Copyright (c) 2009-2010, Warwick Warp Limited
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the Warwick Warp Limited nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL WARWICK WARP LIMITED BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#ifndef WARRAY_IMPL_WARRAY_DFT_H_INCLUDED
#define WARRAY_IMPL_WARRAY_DFT_H_INCLUDED

#include <complex>
#include <cmath>
#include "../warray.h"

namespace warp {
namespace detail {
namespace {
const double pi=3.14159265358979323846;
} // namespace


//! \brief Naive implementation of a 1D DFT
template<typename real_type>
class dif_fft_1d_implementor
{
	private:

		typedef std::complex<real_type>       complex_type;
		typedef warp::arraynd<complex_type,1> complex_buff_type;
		typedef warp::arraynd<size_t,1>       index_buff_type;

	public:

	//! \brief Initialise a DFT for n elements
	dif_fft_1d_implementor(size_t N, direction_e dir, scramble_e scramble) :
		m_buffer(warp::extents[N]),
		m_xl(warp::extents[N]),
		m_num_levels(0),
		m_direction(dir),
		m_scramble(scramble)
	{
		std::bitset<8*sizeof(size_t)> dim_bits(N);

		// check that we have been passed a sequence of a length that is a power of 2
		assert( dim_bits.count() == 1 );

		// crawl along the bit vector until we find the bit that is set, this
		// should gives us log_2(n), which is the number of levels of decomposition
		for( m_num_levels = 0; i < dim_bits.size(); ++m_num_levels )
			if( dim_bits[m_num_levels] )
				break;

		compute_twiddle_factors();

		compute_reordering();
	}


private:

	void compute_twiddle_factors()
	{
		const size_t N = m_buffer.dim(dim1)

		// work out how many twiddle factors there are, this is
		// \sum_{l=1}^{L-1} \frac{N}{2^l}
		// where L is the number of levels. We use the formula for a 
		// geometric sum to get the elegant expression below (which 
		// requires some fiddling to get to it)
		size_t twiddle_dim = ((N << m_num_levels) - N) >> m_num_levels;

		m_twiddles = complex_buff_type(warp::extents[twiddle_dim]);

		// number of points in each block
		size_t Nb = N; 

		// offset into twiddle array
		size_t offset = 0;

		// for each level of the tree
		for( size_t l = 0; l < m_num_levels; ++l )
		{
			size_t half = Nb >> 1;

			// for each value of n that is required at this level
			for( size_t n = 0; n < half; ++n )
			{
				// W_Nb^n = exp(-j2\pi n/N)
				real_type theta = real_type(2.0f) * real_type(pi) * n / real_type(Nb);
				real_type x     = std::cos(theta);
				real_type y     = std::sin(theta);

				m_twiddles[offset+n] = complex_type(x, m_direction == FORWARD ? -y : y);
			}

			offset += half;

			Nb = Nb >> 1;
		}
	}
	
	// Makes a pass through the computation tree and indexes into the indexes of the outputs
	// that the previous level will compute, e.g.
	//       {0,1,2,3,4,5,6,7}
	//   {0,2,4,6}         {1,3,5,7}
	// {0,4}    {2,6}      {1,5} {3,7}
	// {0} {4} {2} {6}    {1} {5} {3} {7}
	void compute_reordering()
	{
		const size_t N = m_buffer.dim(dim1);

		// number of blocks at each level (only 1 at level 0)
		size_t Bl = 1;

		// number of points in each block at current level (N for level 0)
		size_t Nb = N;

		// the ordering of the outputs at the top level of the tree
		index_buff_type indices(warp::extents[N]);
		for( size_t n = 0; n < N; ++n )
			indices[n] = n;

		// for each level of the tree
		for( size_t level = 0; level < m_num_levels; ++level )
		{
			size_t half = Nb / 2;

			// the index variable to the output
			size_t r  = 0;

			// offset to where the current block's data begins
			size_t bo = 0;

			// for each block at this level
			for( size_t b = 0; b < Bl; ++b )
			{
				for( size_t n = 0; n < half; ++n, ++r )
				{
					m_ordering[bo+n]      = indices[2*r];
					m_ordering[bo+n+half] = indices[2*r+1];
				}

				bo += Nb;
			}

			// on the next pass, we will index into these indices
			indices = m_ordering;

			Bl = Bl << 1;
			Nb = Nb >> 1;
		}

	}

	//! \brief Compute the 1D FFT
	real_type operator() (const complex_type* in, int in_stride, complex_type* out, int out_stride)
	{
		const size_t N = m_dim;

		// if forward, we don't need to reorder the coefficients
		if( FORWARD == m_direction )
		{
			for( size_t n = 0; n < N; ++n )
				m_buffer[n] = in[n * in_stride];
		}
		// if backward, the coefficients need to be reordered
		else
		{
			if( UNSCRAMBLED == m_scamble )
			{
				for( size_t n = 0; n < N; ++n )
					m_buffer[n] = in[m_ordering[n] * in_stride];
			}
			// input is already scrambled
			else
			{
				for( size_t n = 0; n < N; ++n )
					m_buffer[n] = in[n * in_stride];
			}
		}


		// number of blocks at each level (only 1 at level 0)
		size_t Bl = 1;

		// number of points in each block at current level (N for level 0)
		size_t Nb = N;

		// the offset into the twiddle factor array for the current level
		size_t offset = 0;

		// for each level of the tree
		for( size_t level = 0; level < m_num_levels; ++level )
		{
			size_t half = Nb / 2;

			// offset to where the current block's data begins
			size_t bo = 0;

			// the input to the FFTs at each level (initially, the input data)
			m_xl.inject(m_buffer);

			// for each block at this level
			for( size_t b = 0; b < Bl; ++b )
			{
				for( size_t n = 0; n < half; ++n )
				{
					complex_t xe =  m_xl[bo+n] + m_xl[bo+n+half];
					complex_t xo = (m_xl[bo+n] - m_xl[bo+n+half]) * m_twiddles[offset+n];

					m_buffer[bo+n]      = xe;
					m_buffer[bo+n+half] = xo;
				}

				bo += Nb;
			}

			offset += half;
			Bl      = Bl << 1;
			Nb      = Nb >> 1;
		}

		// pull the output data out of the buffer and reorder it 
		real_type scale = real_type(1.0) / std::sqrt(real_type(N));

		if( UNSCRAMBLED == m_scramble )
		{
			for( unsigned k = 0; k < N; ++k )
				out[k*out_stride] = scale * m_buffer[m_ordering[k]];
		}
		// output may be scrambled
		else
		{
			for( unsigned k = 0; k < N; ++k )
				out[k*out_stride] = scale * m_buffer[k];
		}

		// this version of the algorithm does not scale
		return real_type(1.0);
	}
	
	//! \brief Determine the size of DFT computed
	size_t dim() const
	{
		return m_coeffs.dim(0);
	}
	
private:

	//! Input buffer used for calculations and packing of variable-stride input.
	complex_buff_type m_buffer;
	
	//! Working buffer used at each level of the tree.
	complex_buff_type m_xl;

	//! The number of levels of the tree.
	size_t m_num_levels;

	//! The twiddle factors.
	complex_buff_type m_twiddles;

	//! The mapping from indices in the output to normal-order coefficients.
	index_buff_type m_ordering;

	//! Direction of transform (forward or backward).
	const direction_e m_direction;

	//! Can the output be in a funny order? i.e. we don't need to do the reordering
	const scramble_e m_scramble;
};



//! \brief Naive implementation of an N-dimensional DFT
template<typename T, size_t N>
class dif_fft_nd_implementor
{
private:
	typedef std::complex<T> complex_type;

public:
	//! \brief 
	dif_fft_nd_implementor(const warp::array_extents<N>& dims, direction_e dir, scramble_e scramble) :
		m_n_minus_one(dims.slice_one(), dir, scramble),
		m_implementor(dims.dim(0))
	{
	}
	
	//! \brief Compute an N dimensional FFT
	template<typename Tptr>
	real_t operator() (const const_arraynd<complex_type,N,Tptr>& in, const arraynd<std::complex<T>,N>& out)
	{
		bounds_check(in,out);
#ifdef BOUNDS_CHECK
		assert(in.dim(warp::dim1) == m_implementor.dim());
#endif
		real_t total_scale = 0;

		// transform all dimensions except this one (the final, or (N-1)th)
		for(size_t i=0;i<m_implementor.dim();++i)
		{
			total_scale += m_n_minus_one(in.slice(warp::indices[i]), out.slice(warp::indices[i]));
		}
		
		// transform the final dimension
		size_t M = in.size()/in.dim(warp::dim1);
		complex_type* optr = out.c_array();
		complex_type* oend = optr+M;
		for(;optr!=oend;++optr)
			total_scale += m_implementor.naive_dft_fwd(optr,out.stride(warp::dim1));
	}

protected:
	dif_fft_nd_implementor<T,N-1> m_n_minus_one;
	dif_fft_1d_implementor<T>     m_implementor;
};

// Specialisation for the one dimensional case
template<typename real_type>
class dif_fft_nd_implementor<real_type,1>
{
	private:

		typedef std::complex<real_type> complex_type;

public:
	dif_fft_nd_implementor(const warp::array_extents<1>& dims, direction_e dir, scramble_e scramble) :
		m_implementor(dims.dim(0), dir, scramble)
	{
	}
	
	template<typename Tptr>
	real_t operator() (const const_arraynd<complex_type,1,Tptr>& in, const arraynd<complex_type,1>& out)
	{
		bounds_check(in,out);
#ifdef BOUNDS_CHECK
		assert(in.dim(warp::dim1) == m_implementor.dim());
#endif
		m_implementor(in.c_array(), in.stride(warp::dim1), out.c_array(), out.stride(warp::dim1));
	}
	
protected:
	dif_fft_1d_implementor<real_type> m_implementor;
};

} // namespace detail
} // namespace warp

#endif

